package com.googlecode.websecuritychecks;

import java.io.IOException;
import java.util.UUID;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import com.googlecode.websecuritychecks.model.*;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;

/**
 * 
 * @author Sai Kris
 *
 */
public class CSRFCheck implements Check {

    private CheckFailureCallback checkFailureCallback;

    @Override
    public void check(final Element ruleXmlElementFragment, final HttpServletRequest httpServletRequest, final HttpServletResponse httpServletResponse) {
        NodeList rule = ruleXmlElementFragment.getElementsByTagNameNS("http://sai.com/websecuritychecks", "CsrfCheck");
        if (rule.getLength() == 1) {
            String description = rule.item(0).getAttributes().getNamedItem("description").getNodeValue();
            String failureStatusCode = rule.item(0).getAttributes().getNamedItem("failureStatusCode").getNodeValue();
            String csrfTokenParamName = rule.item(0).getAttributes().getNamedItem("csrfTokenParamName").getNodeValue();

            String csrfTokenInSession = (String) httpServletRequest.getSession().getAttribute(csrfTokenParamName);

            // If not present in the session, generate it for the first time.
            if (csrfTokenInSession == null) {
                String token = UUID.randomUUID().toString() + "-" + UUID.randomUUID().toString();
                httpServletRequest.getSession().setAttribute(csrfTokenParamName, token);
                httpServletRequest.setAttribute(csrfTokenParamName, token);
            }

            else {
                for (int i = 0; i < rule.item(0).getChildNodes().getLength(); i++) {
					if(httpServletResponse.isCommitted() == true) {
                        break;
					}
                    String urlPattern = rule.item(0).getChildNodes().item(i).getAttributes().getNamedItem("value").getNodeValue();
                    String requestMethod = rule.item(0).getChildNodes().item(i).getAttributes().getNamedItem("method").getNodeValue();

                    if (httpServletRequest.getRequestURI().contains(urlPattern.replace("*", "")) && httpServletRequest.getMethod().equalsIgnoreCase(requestMethod)) {
                        if (csrfTokenInSession.equals(httpServletRequest.getParameter(csrfTokenParamName)) == false) {
                            try {
                                checkFailureCallback.onFailure(httpServletRequest, httpServletResponse, new CheckFailureMetadata(rule.item(0).getNodeName(), description, null,
                                        null, null));
                                if (httpServletResponse.isCommitted() == false) {
                                    httpServletResponse.sendError(Integer.parseInt(failureStatusCode), "CSRF Attack identified.");
                                }
                            } catch (NumberFormatException | IOException e) {
                                e.printStackTrace();
                            }
                        }
                    }

                }
            }

        }

    }

    @Override
    public void addCheckFailureCallback(final CheckFailureCallback checkFailureCallback) {
        this.checkFailureCallback = checkFailureCallback;
    }

}
